/*
 * Copyright (c) 2018-2024 北京华恒盛世科技有限公司
 * QUICKPOOL is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *     http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
*/

#include "jm.h"

#define CGROUP_SUBSYSTEM_MEMORY "memory"
#define CGROUP_SUBSYSTEM_CPUACCT "cpuacct"

#define MAX_UINT64_STR_LENGTH 21

static int ticks_per_second(void);
static int cgroup_get_job_cputime(JOBID_T, uint64_t *, uint64_t *);
static int cgroup_get_job_mem_usage(JOBID_T, uint64_t *, uint64_t *);
static int cgroup_get_job_pids(JOBID_T, int *, int **);
static void rlimit_memory(struct resourceLimit *, struct rlimit *);

/**
 * @brief       Setup job cgroup before starting job.
 *
 * @param[in]   jobid        #1: job id
 * @param[in]   jobpid       #2: job pid
 * @param[in]   subsystem    #3: name of cgroup subsystem
 *
 * @retval      0            #1: succeed
 * @retval      -1           #2: failed, syscall fail.
 *
 * @note        Need call this function to setup job cgroup.
 *              We support 2 cgroup subsystems so far. One is memory, the other is cpuacct
 */
int
cgroup_setup_job_entry(JOBID_T jobid, pid_t jobpid, const char *subsystem)
{
    char *strCgroupRoot;
    strCgroupRoot = baseParams[OL_CGROUP_ROOT].string_value;
    if (NULL == strCgroupRoot) {
        return 0;
    }

    char strJobId[32];
    jobid_2string(jobid, strJobId);

    if (mkdir(strCgroupRoot, 0755) < 0 && errno != EEXIST) {
        log_message(ERR, NOCLASS, "%s: mkdir(%s) failed when set job <%s> cgroup subsystem %s due to %m.",
                    __func__, strCgroupRoot, strJobId, subsystem);
        return -1;
    }

    char strSubSystem[PATH_MAX+1];
    snprintf(strSubSystem, PATH_MAX, "%s/%s", strCgroupRoot, subsystem);
    if (mkdir(strSubSystem, 0755) < 0 && errno != EEXIST) {
        log_message(ERR, NOCLASS, "%s: mkdir(%s) failed when set job <%s> cgroup due to %m.",
                    __func__, strSubSystem, strJobId);
        return -1;
    }

    snprintf(strSubSystem, PATH_MAX, "%s/%s/quickpool", strCgroupRoot, subsystem);
    if (mkdir(strSubSystem, 0755) < 0 && errno != EEXIST) {
        log_message(ERR, NOCLASS, "%s: mkdir(%s) failed when set job <%s> cgroup due to %m.",
                    __func__, strSubSystem, strJobId);
        return -1;
    }

    snprintf(strSubSystem, PATH_MAX, "%s/%s/quickpool/job.%s", strCgroupRoot, subsystem, strJobId);
    if (mkdir(strSubSystem, 0755) < 0 && errno != EEXIST) {
        log_message(ERR, NOCLASS, "%s: mkdir(%s) failed when set job cgroup due to %m.",
                    __func__, strSubSystem);
        return -1;
    }
    
    FILE *fp;
    snprintf(strSubSystem, PATH_MAX, "%s/%s/quickpool/job.%s/tasks",
             strCgroupRoot, subsystem, strJobId);
    fp = fopen(strSubSystem, "a");
    if (NULL == fp) {
        log_message(ERR, NOCLASS, "%s: fopen(%s) failed when set job cgroup due to %m.",
                    __func__, strSubSystem);
        return -1;
    }

    if (fprintf(fp, "%d\n", jobpid) < 0) {
        log_message(ERR, NOCLASS, "%s: fprintf(%s) failed when set job cgroup due to %m.",
                    __func__, strSubSystem);
        fclose(fp);
        return -1;
    }
    fclose(fp);
    
    return 0;
} // end function cgroup_setup_job_entry

/**
 * @brief       Set job memory limit through cgroup.
 *
 * @param[in]   job          #1: job entry
 *
 * @note        Write job memory limit to file memory.limit_in_bytes
 */
void
setup_mem_cgroup(struct jobEntry *job)
{
    struct rlimit rlimit;
    rlimit_memory(&job->job_base.resource_limits[RESOURCE_LIMIT_RSS], &rlimit);

    char *strCgroupRoot;
    strCgroupRoot = baseParams[OL_CGROUP_ROOT].string_value;
    if (NULL == strCgroupRoot) {
        return;
    }

    char strJobId[32];
    jobid_2string(job->job_base.jobid, strJobId);

    char strCgroupMemory[PATH_MAX+1];
    snprintf(strCgroupMemory, PATH_MAX, "%s/memory/quickpool/job.%s/memory.limit_in_bytes",
              strCgroupRoot, strJobId);

    FILE *fp;
    fp = fopen(strCgroupMemory, "a+");
    if (NULL == fp) {
        log_message(ERR, NOCLASS, "%s: fopen(%s) failed when set memory limit for job due to %m.",
                    __func__, strCgroupMemory);
        return;
    }

    if (fprintf(fp, "%lu\n", rlimit.rlim_cur) < 0) {
        log_message(ERR, NOCLASS, "%s: fprintf(%s) failed when set memory limit for job due to %m.",
                    __func__, strCgroupMemory);
        fclose(fp);
        return;
    }
    fclose(fp);

    log_message(DEBUG, EXECUTE, "%s: job <%s> memory limit is set. cur %lu max %lu",
                __func__, strJobId, rlimit.rlim_cur, rlimit.rlim_max);

    return;
} // end function setup_mem_cgroup

/**
 * @brief       Remove job cgroup after job is finished.
 *
 * @param[in]   jobid        #1: job id
 * @param[in]   subsystem    #2: name of cgroup subsystem
 *
 * @retval      0            #1: succeed
 * @retval      -1           #2: failed, syscall fail.
 */
int
cgroup_rm_job_entry(JOBID_T jobid, const char *subsystem)
{
    char strJobId[32];
    jobid_2string(jobid, strJobId);

    char *strCgroupRoot;
    strCgroupRoot = baseParams[OL_CGROUP_ROOT].string_value;
    if (NULL == strCgroupRoot) {
        return 0;
    }

    char strCgroupJob[PATH_MAX+1];
    snprintf(strCgroupJob, PATH_MAX, "%s/%s/quickpool/job.%s", strCgroupRoot, subsystem, strJobId);
    if (rmdir(strCgroupJob) < 0 && errno != ENOENT) {
        log_message(ERR, NOCLASS, "%s: rmdir(%s) failed when remove job cgroup due to %m.",
                    __func__, strCgroupJob);
        return -1;
    }

    return 0;
} //end function cgroup_rm_job_entry

/**
 * @brief       Get job resource usage from job cgroup files.
 *
 * @param[in]   jobid        #1: job id
 *
 * @retval      not NULL     #1: succeed
 * @retval      NULL         #2: failed, syscall fail.
 *
 * @note        Read job cgroup files to get job cpu time, memory usage and pids.
 */
struct jobResourceUsage *
cgroup_get_job_rusage(JOBID_T jobid)
{
    struct jobResourceUsage *jRusage;
    jRusage = (struct jobResourceUsage *)calloc(1, sizeof(struct jobResourceUsage));
    if (NULL == jRusage) {
        return NULL;
    }

    uint64_t utime=0, stime=0;
    if (cgroup_get_job_cputime(jobid, &utime, &stime) < 0) {
        utime = 0;
        stime = 0;
    }

    uint64_t mem=0, swap=0;
    if (cgroup_get_job_mem_usage(jobid, &mem, &swap) < 0) {
        mem = 0;
        swap = 0;
    }

    int npid = 0;
    int *pids = NULL;
    if (cgroup_get_job_pids(jobid, &npid, &pids) < 0) {
        npid = 0;
        pids = NULL;
    }

    if (0 < npid) {

        struct pidInfo *info;
        info = (struct pidInfo *)calloc(npid, sizeof(struct pidInfo));
        if (NULL == info) {
            FREEUP(jRusage);
            FREEUP(pids);
            return NULL;
        }

        int *pgids;
        pgids = (int *)calloc(npid, sizeof(int));
        if (NULL == pgids) {
            FREEUP(jRusage);
            FREEUP(pids);
            FREEUP(info);
            return NULL;
        }

        int npgid = 0, i, j;
        for (i = 0; i < npid; i++) {
            int pgid = getpgid(pids[i]);
            
            info[i].pid = pids[i];
            info[i].pgid = pgid;

            for (j = 0; j < npgid; j++) {
                if (pgids[j] == pgid) {
                    break;
                }
            }

            if (j == npgid) {
                pgids[npgid++] = pgid;
            }
        }

        FREEUP(pids);

        if (npgid == 0) {
            FREEUP(pgids);
        } else {
            int *newPGids;
            newPGids = (int *)realloc(pgids, npgid * sizeof(int));
            if (NULL == newPGids) {
                FREEUP(pgids);
                FREEUP(jRusage);
                FREEUP(info);
                return NULL;
            }
            pgids = newPGids;
        }

        jRusage->npgids = npgid;
        jRusage->pgid= pgids;
        jRusage->npids = npid;
        jRusage->pidinfo = info;
    } else {
        FREEUP(pids);
    }

    jRusage->utime = (int)utime;
    jRusage->stime = (int)stime;
    jRusage->mem = (int)(mem / 1024);
    jRusage->swap = (int)(swap / 1024);
    
    return jRusage;
} // end function cgroup_get_job_rusage

/* ------------------------------------------------
 *
 *  below are static functions used in this file
 *
 * ------------------------------------------------ */

static int 
ticks_per_second(void)
{
    int ticks;

    ticks = sysconf(_SC_CLK_TCK);

    return ticks;
}

/**
 * @brief       Get job cputime from cgroup file cpuacct.stat.
 *
 * @param[in]   jobid        #1: job id
 * @param[out]  utime        #2: job cpu user time
 * @param[out]  stime        #3: job cpu system time
 *
 * @retval      0            #1: succeed
 * @retval      -1           #2: failed
 */
static int 
cgroup_get_job_cputime(JOBID_T jobid, uint64_t *utime, uint64_t *stime)
{
    if (NULL == utime && NULL == stime) {
        return -1;
    }

    char *strCgroupRoot;
    strCgroupRoot = baseParams[OL_CGROUP_ROOT].string_value;
    if (NULL == strCgroupRoot) {
        return 0;
    }

    char strJobId[32];    
    jobid_2string(jobid, strJobId);

    char strCgroupJob[PATH_MAX+1];
    snprintf(strCgroupJob, PATH_MAX, "%s/cpuacct/quickpool/job.%s/cpuacct.stat",
             strCgroupRoot, strJobId);

    FILE *fp;
    fp = fopen(strCgroupJob, "r");
    if (NULL == fp) {
        return -1;
    }

    int cc;
    uint64_t ut, st;
    cc = fscanf(fp, "%*s%llu%*s%llu", (long long unsigned int *)&ut, (long long unsigned int *)&st);
    if (2 != cc) {
        fclose(fp);
        return -1;
    }
    fclose(fp);

    if (NULL != utime) {
        *utime = ut / (uint64_t) ticks_per_second();
    }

    if (NULL != stime) {
        *stime = st / (uint64_t) ticks_per_second();
    }
       
    return 0;
} // end function cgroup_get_job_cputime

/**
 * @brief       Get job memory and swap usage from cgroup files:
 *              memory.usage_in_bytes and memory.memsw.usage_in_bytes.
 *
 * @param[in]   jobid        #1: job id
 * @param[out]  outmem       #2: job used memory
 * @param[out]  outswap      #3: job used swap
 *
 * @retval      0            #1: succeed
 * @retval      -1           #2: failed
 */
static int 
cgroup_get_job_mem_usage(JOBID_T jobid, uint64_t *outmem, uint64_t *outswap)
{
    if (NULL == outmem && NULL == outswap) {
        return -1;
    }

    char *strCgroupRoot;
    strCgroupRoot = baseParams[OL_CGROUP_ROOT].string_value;
    if (NULL == strCgroupRoot) {
        return 0;
    }
    
    char strJobId[32];
    jobid_2string(jobid, strJobId);

    char strCgroupMem[PATH_MAX+1];
    snprintf(strCgroupMem, PATH_MAX, "%s/memory/quickpool/job.%s/memory.usage_in_bytes",
             strCgroupRoot, strJobId);

    FILE *fp;
    fp = fopen(strCgroupMem, "r");
    if (NULL == fp) {
        return -1;
    }

    int cc;
    uint64_t mem, memsw;
    cc = fscanf(fp, "%llu", (long long unsigned int *)&mem);
    if (1 != cc) {
        fclose(fp);
        return -1;
    }
    fclose(fp);

    snprintf(strCgroupMem, PATH_MAX, "%s/memory/quickpool/job.%s/memory.memsw.usage_in_bytes",
             strCgroupRoot, strJobId);
    fp = fopen(strCgroupMem, "r");
    if (NULL == fp) {
        return -1;
    }

    cc = fscanf(fp, "%llu", (long long unsigned int *)&memsw);
    if (1 != cc) {
        fclose(fp);
        return -1;
    }
    fclose(fp);

    if (NULL != outmem) {
        *outmem = mem;
    }

    if (NULL != outswap) {
        if (memsw > mem) {
            *outswap = memsw - mem;
        } else {
            *outswap = 0;
        }
    }

    return 0;
} // end function cgroup_get_job_mem_usage

/**
 * @brief       Get job pids from cgroup file cgroup.procs
 *
 * @param[in]   jobid        #1: job id
 * @param[out]  outnpid      #2: number of job pids
 * @param[out]  outpids      #3: job pids
 *
 * @retval      0            #1: succeed
 * @retval      -1           #2: failed
 */
static int 
cgroup_get_job_pids(JOBID_T jobid, int *outnpid, int **outpids)
{
    if (NULL == outnpid || NULL == outpids) {
        return -1;
    }

    *outnpid = 0;
    *outpids = NULL;

    char *strCgroupRoot;
    strCgroupRoot = baseParams[OL_CGROUP_ROOT].string_value;
    if (NULL == strCgroupRoot) {
        return 0;
    }

    int size=32;
    int *pids;
    pids = (int *)calloc(size, sizeof(int));
    if (NULL == pids) {
        return -1;
    }
    
    char strJobId[32];
    jobid_2string(jobid, strJobId);

    char strCgroupProcs[PATH_MAX+1];
    snprintf(strCgroupProcs, PATH_MAX, "%s/memory/quickpool/job.%s/cgroup.procs",
             strCgroupRoot, strJobId);

    FILE *fp;
    fp = fopen(strCgroupProcs, "r");
    if (NULL == fp) {
        FREEUP(pids);
        return -1;
    }

    int npid=0, pid;
    while(fscanf(fp, "%d", &pid) != EOF) {

        pids[npid++] = pid;

        if (npid == size) {

            int *newPids;
            size *= 2;
            newPids = (int *)realloc(pids, size * sizeof(int));
            if (NULL == newPids) {
                FREEUP(pids);
                fclose(fp);
                return -1;
            }
            pids = newPids;
        }
    }

    fclose(fp);

    *outnpid = npid;
    *outpids = pids;

    return 0;
} // end function cgroup_get_job_pids

static void
rlimit_memory(struct resourceLimit *reslimit, struct rlimit *rlimit)
{
    int softok = FALSE;
    int hardok = FALSE;

    if (reslimit->rlim_curl == 0xffffffff && reslimit->rlim_curh == 0x7fffffff) {
        rlimit->rlim_cur = RLIM_INFINITY;
        softok = TRUE;
    }

    if (reslimit->rlim_maxl == 0xffffffff && reslimit->rlim_maxh == 0x7fffffff) {
        rlimit->rlim_max = RLIM_INFINITY;
        hardok = TRUE;
    }

    if (sizeof(rlimit->rlim_cur) == 8) {
        if (!softok) {
            rlimit->rlim_cur = reslimit->rlim_curh;
            rlimit->rlim_cur = rlimit->rlim_cur << 32;
            rlimit->rlim_cur |= (reslimit->rlim_curl & 0xffffffff);
        }

        if (!hardok) {
            rlimit->rlim_max = reslimit->rlim_maxh;
            rlimit->rlim_max = rlimit->rlim_max << 32;
            rlimit->rlim_max |= (reslimit->rlim_maxl & 0xffffffff);
        }
    } else {
        if (!softok) {
            if ((reslimit->rlim_curh > 0) || (reslimit->rlim_curl & 0x80000000)) {
                rlimit->rlim_cur = RLIM_INFINITY;
            } else {
                rlimit->rlim_cur = reslimit->rlim_curl;
            }
        }

        if (!hardok) {
            if ((reslimit->rlim_maxh > 0) || (reslimit->rlim_maxl & 0x80000000)) {
                rlimit->rlim_max = RLIM_INFINITY;
            } else {
                rlimit->rlim_max = reslimit->rlim_maxl;
            }
        }
    }
} // end function rlimit_memory
